import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict, Counter
import datetime
import re
import argparse

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class StockAnalyzer:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.stock_data = {}
        self.stock_frequency = Counter()
        self.stock_info = {}
        self.load_data()
        
    def load_data(self):
        """加载所有股票数据文件"""
        print(f"正在从 {self.data_dir} 加载数据...")
        files = [f for f in os.listdir(self.data_dir) if f.endswith('.txt')]
        files.sort()  # 按日期排序
        
        for file in files:
            date_str = file.split('.')[0]
            file_path = os.path.join(self.data_dir, file)
            try:
                # 读取文件数据
                df = pd.read_csv(file_path, sep='\t')
                # 记录每只股票出现的次数
                for _, row in df.iterrows():
                    stock_code = row['代码']
                    self.stock_frequency[stock_code] += 1
                    
                    # 保存股票信息
                    if stock_code not in self.stock_info:
                        self.stock_info[stock_code] = {
                            '名称': row['名称'],
                            '行业': '未知',  # 数据中没有行业信息，可以后续添加
                            '最早出现日期': date_str,
                            '最新出现日期': date_str,
                            '出现次数': 1,
                            '平均涨跌幅': row['涨跌幅'],
                            '累计涨跌幅': row['涨跌幅'],
                            '最大涨幅': row['涨跌幅'],
                            '最小涨幅': row['涨跌幅'],
                            '平均换手率': row['换手率'] if '换手率' in row else 0,
                            '平均成交量': row['成交量'],
                            '市盈率': row['市盈率-动态'] if '市盈率-动态' in row else 0,
                            '市净率': row['市净率'] if '市净率' in row else 0,
                            '60日涨跌幅': row['60日涨跌幅'] if '60日涨跌幅' in row else 0,
                            '年初至今涨跌幅': row['年初至今涨跌幅'] if '年初至今涨跌幅' in row else 0,
                            '数据': [{
                                '日期': date_str,
                                '价格': row['最新价'],
                                '涨跌幅': row['涨跌幅'],
                                '成交量': row['成交量'],
                                '换手率': row['换手率'] if '换手率' in row else 0
                            }]
                        }
                    else:
                        # 更新股票信息
                        info = self.stock_info[stock_code]
                        info['最新出现日期'] = date_str
                        info['出现次数'] += 1
                        info['累计涨跌幅'] += row['涨跌幅']
                        info['平均涨跌幅'] = info['累计涨跌幅'] / info['出现次数']
                        info['最大涨幅'] = float(max(float(info['最大涨幅']), float(row['涨跌幅'])))
                        info['最小涨幅'] = float(min(float(info['最小涨幅']), float(row['涨跌幅'])))
                        info['平均换手率'] = (info['平均换手率'] * (info['出现次数'] - 1) + 
                                         (row['换手率'] if '换手率' in row else 0)) / info['出现次数']
                        info['平均成交量'] = (info['平均成交量'] * (info['出现次数'] - 1) + row['成交量']) / info['出现次数']
                        info['市盈率'] = row['市盈率-动态'] if '市盈率-动态' in row else info['市盈率']
                        info['市净率'] = row['市净率'] if '市净率' in row else info['市净率']
                        info['60日涨跌幅'] = row['60日涨跌幅'] if '60日涨跌幅' in row else info['60日涨跌幅']
                        info['年初至今涨跌幅'] = row['年初至今涨跌幅'] if '年初至今涨跌幅' in row else info['年初至今涨跌幅']
                        info['数据'].append({
                            '日期': date_str,
                            '价格': row['最新价'],
                            '涨跌幅': row['涨跌幅'],
                            '成交量': row['成交量'],
                            '换手率': row['换手率'] if '换手率' in row else 0
                        })
                
                # 保存每日数据
                self.stock_data[date_str] = df
                
            except Exception as e:
                print(f"处理文件 {file} 时出错: {e}")
        
        print(f"共加载 {len(files)} 个交易日的数据，包含 {len(self.stock_info)} 只股票")
    
    def get_frequent_stocks(self, min_frequency=1):
        """获取出现频率较高的股票"""
        return [code for code, freq in self.stock_frequency.items() if freq >= min_frequency]
    
    def calculate_stability_score(self, stock_code):
        """计算股票稳定性得分"""
        if stock_code not in self.stock_info:
            return 0
        
        info = self.stock_info[stock_code]
        if info['出现次数'] < 2:
            return 0
        
        # 计算涨跌幅的标准差
        price_changes = [data['涨跌幅'] for data in info['数据']]
        std_dev = np.std(price_changes)
        
        # 稳定性得分 = 平均涨跌幅 / (标准差 + 1)，加1是为了避免除以0
        stability = info['平均涨跌幅'] / (std_dev + 1)
        
        return stability
    
    def calculate_momentum_score(self, stock_code):
        """计算股票动量得分"""
        if stock_code not in self.stock_info:
            return 0
        
        info = self.stock_info[stock_code]
        if info['出现次数'] < 2:
            return 0
        
        # 最近的涨跌幅权重更高
        recent_changes = [data['涨跌幅'] for data in info['数据'][-3:]]
        if not recent_changes:
            return 0
        
        # 动量得分 = 最近涨跌幅的加权平均
        weights = np.array([0.5, 0.3, 0.2])[:len(recent_changes)]
        weights = weights / weights.sum()  # 归一化权重
        momentum = np.average(recent_changes, weights=weights)
        
        return momentum
    
    def calculate_volume_score(self, stock_code):
        """计算成交量得分"""
        if stock_code not in self.stock_info:
            return 0
        
        info = self.stock_info[stock_code]
        if info['出现次数'] < 2:
            return 0
        
        # 计算最近成交量变化
        volumes = [data['成交量'] for data in info['数据']]
        if len(volumes) < 2:
            return 0
        
        # 成交量增长率
        volume_growth = (volumes[-1] / volumes[0] - 1) if volumes[0] > 0 else 0
        
        # 成交量得分 = 成交量增长率 * 平均换手率
        volume_score = volume_growth * info['平均换手率'] / 10  # 除以10进行归一化
        
        return volume_score
    
    def calculate_pe_score(self, stock_code):
        """计算市盈率得分（市盈率越低越好）"""
        if stock_code not in self.stock_info:
            return 0
        
        info = self.stock_info[stock_code]
        pe = info['市盈率']
        
        # 负的市盈率或者异常高的市盈率都不好
        if pe <= 0 or pe > 100:
            return 0
        
        # 市盈率得分 = 1 / 市盈率 * 10（归一化）
        pe_score = 1 / pe * 10
        
        return pe_score
    
    def calculate_pb_score(self, stock_code):
        """计算市净率得分（市净率越低越好）"""
        if stock_code not in self.stock_info:
            return 0
        
        info = self.stock_info[stock_code]
        pb = info['市净率']
        
        # 负的市净率或者异常高的市净率都不好
        if pb <= 0 or pb > 10:
            return 0
        
        # 市净率得分 = 1 / 市净率 * 5（归一化）
        pb_score = 1 / pb * 5
        
        return pb_score
    
    def calculate_comprehensive_score(self, stock_code, weights=None):
        """计算综合得分，可自定义各指标权重"""
        if stock_code not in self.stock_info:
            return 0
        
        # 默认权重
        if weights is None:
            weights = {
                'stability': 0.3,  # 稳定性
                'momentum': 0.4,   # 动量
                'volume': 0.2,     # 成交量
                'frequency': 0.1,  # 出现频率
                'pe': 0.0,         # 市盈率
                'pb': 0.0          # 市净率
            }
        
        info = self.stock_info[stock_code]
        
        # 各项指标得分
        stability = self.calculate_stability_score(stock_code) * weights['stability']
        momentum = self.calculate_momentum_score(stock_code) * weights['momentum']
        volume = self.calculate_volume_score(stock_code) * weights['volume']
        frequency = (info['出现次数'] / 10) * weights['frequency']  # 归一化频率得分
        pe = self.calculate_pe_score(stock_code) * weights['pe']
        pb = self.calculate_pb_score(stock_code) * weights['pb']
        
        # 综合得分
        score = stability + momentum + volume + frequency + pe + pb
        
        return score
    
    def get_top_stocks(self, n=10, min_frequency=1, weights=None):
        """获取综合评分最高的前N只股票"""
        frequent_stocks = self.get_frequent_stocks(min_frequency)
        
        # 计算每只股票的综合得分
        stock_scores = [(code, self.calculate_comprehensive_score(code, weights)) for code in frequent_stocks]
        
        # 按得分降序排序
        stock_scores.sort(key=lambda x: x[1], reverse=True)
        
        return stock_scores[:n]
    
    def get_stock_recommendation_reasons(self, stock_code):
        """获取股票推荐理由"""
        if stock_code not in self.stock_info:
            return "无法找到该股票的信息"
        
        info = self.stock_info[stock_code]
        reasons = []
        
        # 基于出现频率的理由
        if info['出现次数'] >= 2:
            reasons.append(f"该股票在分析期间多次出现({info['出现次数']}次)，表明具有持续的市场关注度")
        
        # 基于平均涨跌幅的理由
        if info['平均涨跌幅'] > 1:
            reasons.append(f"平均涨幅较高({info['平均涨跌幅']:.2f}%)，表现优于大多数股票")
        
        # 基于稳定性的理由
        stability = self.calculate_stability_score(stock_code)
        if stability > 0.2:
            reasons.append("涨跌幅波动较小，股价表现稳定")
        
        # 基于动量的理由
        momentum = self.calculate_momentum_score(stock_code)
        if momentum > 1:
            reasons.append("近期涨势良好，具有较强的上涨动能")
        
        # 基于成交量的理由
        if info['平均换手率'] > 2:
            reasons.append(f"平均换手率较高({info['平均换手率']:.2f}%)，交易活跃度高")
        
        # 基于长期表现的理由
        if info['60日涨跌幅'] > 5:
            reasons.append(f"60日涨幅达{info['60日涨跌幅']:.2f}%，中期表现优异")
        
        if info['年初至今涨跌幅'] > 10:
            reasons.append(f"年初至今涨幅达{info['年初至今涨跌幅']:.2f}%，长期表现优异")
        
        # 基于估值的理由
        if 0 < info['市盈率'] < 30:
            reasons.append(f"市盈率为{info['市盈率']:.2f}，估值合理")
        
        if 0 < info['市净率'] < 3:
            reasons.append(f"市净率为{info['市净率']:.2f}，估值合理")
        
        # 如果没有找到特别的理由，给出一个通用理由
        if not reasons:
            reasons.append("综合各项指标表现良好，值得关注")
        
        return "\n".join([f"- {reason}" for reason in reasons])
    
    def generate_recommendation_report(self, n=10, min_frequency=1, weights=None, output_file=None):
        """生成股票推荐报告"""
        top_stocks = self.get_top_stocks(n, min_frequency, weights)
        
        report = ["# 股票推荐报告\n"]
        report.append(f"## 分析周期: {list(self.stock_data.keys())[0]} 至 {list(self.stock_data.keys())[-1]}\n")
        
        # 添加权重信息
        if weights:
            report.append("## 自定义分析权重\n")
            for key, value in weights.items():
                report.append(f"- **{key}**: {value:.2f}\n")
            report.append("\n")
        
        report.append(f"## 推荐股票列表 (共{len(top_stocks)}只)\n")
        
        for i, (code, score) in enumerate(top_stocks, 1):
            info = self.stock_info[code]
            report.append(f"### {i}. {info['名称']}({code})\n")
            report.append(f"**综合评分**: {score:.2f}\n")
            report.append(f"**最新价格**: {info['数据'][-1]['价格']}元\n")
            report.append(f"**平均涨跌幅**: {info['平均涨跌幅']:.2f}%\n")
            report.append(f"**60日涨跌幅**: {info['60日涨跌幅']:.2f}%\n")
            report.append(f"**年初至今涨跌幅**: {info['年初至今涨跌幅']:.2f}%\n")
            report.append(f"**市盈率**: {info['市盈率']}\n")
            report.append(f"**市净率**: {info['市净率']}\n")
            report.append(f"**平均换手率**: {info['平均换手率']:.2f}%\n")
            report.append("**推荐理由**:\n")
            report.append(self.get_stock_recommendation_reasons(code) + "\n")
            report.append("---\n")
        
        report.append("## 分析方法说明\n")
        report.append("本报告基于以下指标对股票进行综合评分:\n")
        report.append("1. **稳定性**: 评估股票价格波动的稳定程度\n")
        report.append("2. **动量**: 评估股票近期的上涨趋势和力度\n")
        report.append("3. **交易活跃度**: 评估股票的成交量和换手率\n")
        report.append("4. **出现频率**: 评估股票在分析期间的出现次数\n")
        if weights and (weights['pe'] > 0 or weights['pb'] > 0):
            report.append("5. **估值指标**: 评估股票的市盈率和市净率\n")
        
        # 添加综合得分计算公式
        formula = "综合得分计算公式: "
        if weights:
            terms = []
            if weights['stability'] > 0:
                terms.append(f"稳定性*{weights['stability']:.2f}")
            if weights['momentum'] > 0:
                terms.append(f"动量*{weights['momentum']:.2f}")
            if weights['volume'] > 0:
                terms.append(f"交易活跃度*{weights['volume']:.2f}")
            if weights['frequency'] > 0:
                terms.append(f"频率*{weights['frequency']:.2f}")
            if weights['pe'] > 0:
                terms.append(f"市盈率得分*{weights['pe']:.2f}")
            if weights['pb'] > 0:
                terms.append(f"市净率得分*{weights['pb']:.2f}")
            formula += " + ".join(terms)
        else:
            formula += "稳定性*0.3 + 动量*0.4 + 交易活跃度*0.2 + 频率*0.1"
        
        report.append(f"\n{formula}\n")
        
        report_text = ''.join(report)
        
        # 如果指定了输出文件，则写入文件
        if output_file:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(report_text)
            print(f"推荐报告已生成: {output_file}")
        
        return report_text
    
    def analyze_by_criteria(self, criteria='comprehensive', n=10, min_frequency=1, weights=None):
        """根据不同标准分析股票"""
        frequent_stocks = self.get_frequent_stocks(min_frequency)
        
        if criteria == 'comprehensive':
            # 综合评分
            scores = [(code, self.calculate_comprehensive_score(code, weights)) for code in frequent_stocks]
            title = "综合评分最高的股票"
            ylabel = "综合评分"
        elif criteria == 'stability':
            # 稳定性评分
            scores = [(code, self.calculate_stability_score(code)) for code in frequent_stocks]
            title = "稳定性最高的股票"
            ylabel = "稳定性评分"
        elif criteria == 'momentum':
            # 动量评分
            scores = [(code, self.calculate_momentum_score(code)) for code in frequent_stocks]
            title = "动量最强的股票"
            ylabel = "动量评分"
        elif criteria == 'volume':
            # 成交量评分
            scores = [(code, self.calculate_volume_score(code)) for code in frequent_stocks]
            title = "交易活跃度最高的股票"
            ylabel = "交易活跃度评分"
        elif criteria == 'frequency':
            # 出现频率
            scores = [(code, self.stock_info[code]['出现次数']) for code in frequent_stocks]
            title = "出现频率最高的股票"
            ylabel = "出现次数"
        elif criteria == 'avg_change':
            # 平均涨跌幅
            scores = [(code, self.stock_info[code]['平均涨跌幅']) for code in frequent_stocks]
            title = "平均涨幅最高的股票"
            ylabel = "平均涨跌幅 (%)"
        elif criteria == 'ytd_change':
            # 年初至今涨跌幅
            scores = [(code, self.stock_info[code]['年初至今涨跌幅']) for code in frequent_stocks]
            title = "年初至今涨幅最高的股票"
            ylabel = "年初至今涨跌幅 (%)"
        elif criteria == 'pe':
            # 市盈率（越低越好）
            scores = [(code, self.calculate_pe_score(code)) for code in frequent_stocks]
            title = "市盈率评分最高的股票（市盈率越低越好）"
            ylabel = "市盈率评分"
        elif criteria == 'pb':
            # 市净率（越低越好）
            scores = [(code, self.calculate_pb_score(code)) for code in frequent_stocks]
            title = "市净率评分最高的股票（市净率越低越好）"
            ylabel = "市净率评分"
        else:
            print(f"未知的分析标准: {criteria}")
            return
        
        # 按评分降序排序
        scores.sort(key=lambda x: x[1], reverse=True)
        top_scores = scores[:n]
        
        # 打印详细信息
        print(f"\n{title}:")
        for i, (code, score) in enumerate(top_scores, 1):
            info = self.stock_info[code]
            print(f"{i}. {info['名称']}({code}) - 评分: {score:.2f}, 平均涨跌幅: {info['平均涨跌幅']:.2f}%, "  
                  f"60日涨跌幅: {info['60日涨跌幅']:.2f}%, 年初至今涨跌幅: {info['年初至今涨跌幅']:.2f}%")
        
        return top_scores
    
    def plot_stock_price_trend(self, stock_code):
        """绘制股票价格走势图"""
        if stock_code not in self.stock_info:
            print(f"找不到股票 {stock_code} 的数据")
            return
        
        info = self.stock_info[stock_code]
        data = info['数据']
        
        # 提取日期和价格
        dates = [d['日期'] for d in data]
        prices = [d['价格'] for d in data]
        changes = [d['涨跌幅'] for d in data]
        
        # 创建图表
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), gridspec_kw={'height_ratios': [3, 1]})
        
        # 绘制价格走势
        ax1.plot(dates, prices, 'b-', marker='o')
        ax1.set_title(f"{info['名称']}({stock_code}) 价格走势")
        ax1.set_ylabel('价格 (元)')
        ax1.grid(True)
        
        # 设置x轴标签旋转
        plt.setp(ax1.get_xticklabels(), rotation=45, ha='right')
        
        # 绘制涨跌幅
        bars = ax2.bar(dates, changes, color=['r' if x > 0 else 'g' for x in changes])
        ax2.set_title('日涨跌幅 (%)')
        ax2.set_ylabel('涨跌幅 (%)')
        ax2.grid(True)
        
        # 设置x轴标签旋转
        plt.setp(ax2.get_xticklabels(), rotation=45, ha='right')
        
        plt.tight_layout()
        plt.show()
    
    def plot_top_stocks_comparison(self, n=10, min_frequency=1, weights=None):
        """绘制前N只股票的表现比较"""
        top_stocks = self.get_top_stocks(n, min_frequency, weights)
        
        # 提取股票名称和平均涨跌幅
        names = [f"{self.stock_info[code]['名称']}({code})" for code, _ in top_stocks]
        avg_changes = [self.stock_info[code]['平均涨跌幅'] for code, _ in top_stocks]
        year_to_date = [self.stock_info[code]['年初至今涨跌幅'] for code, _ in top_stocks]
        sixty_days = [self.stock_info[code]['60日涨跌幅'] for code, _ in top_stocks]
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(14, 8))
        
        # 设置柱状图的位置
        x = np.arange(len(names))
        width = 0.25
        
        # 绘制三组柱状图
        rects1 = ax.bar(x - width, avg_changes, width, label='平均涨跌幅')
        rects2 = ax.bar(x, sixty_days, width, label='60日涨跌幅')
        rects3 = ax.bar(x + width, year_to_date, width, label='年初至今涨跌幅')
        
        # 添加标题和标签
        ax.set_title('推荐股票涨跌幅比较')
        ax.set_ylabel('涨跌幅 (%)')
        ax.set_xticks(x)
        ax.set_xticklabels(names, rotation=45, ha='right')
        ax.legend()
        
        # 添加数据标签
        def autolabel(rects):
            for rect in rects:
                height = rect.get_height()
                ax.annotate(f'{height:.2f}%',
                            xy=(rect.get_x() + rect.get_width() / 2, height),
                            xytext=(0, 3),  # 3点垂直偏移
                            textcoords="offset points",
                            ha='center', va='bottom')
        
        autolabel(rects1)
        autolabel(rects2)
        autolabel(rects3)
        
        plt.tight_layout()
        plt.grid(axis='y')
        plt.show()

def parse_args():
    parser = argparse.ArgumentParser(description='股票分析工具')
    parser.add_argument('--data_dir', type=str, default='d:\\Work\\python\\selfTry\\data', help='数据目录路径')
    parser.add_argument('--output', type=str, default='d:\\Work\\AItest\\custom_stock_recommendation.md', help='输出报告文件路径')
    parser.add_argument('--top_n', type=int, default=10, help='推荐股票数量')
    parser.add_argument('--min_freq', type=int, default=1, help='最小出现频率')
    parser.add_argument('--criteria', type=str, default='comprehensive', 
                        choices=['comprehensive', 'stability', 'momentum', 'volume', 'frequency', 'avg_change', 'ytd_change', 'pe', 'pb'],
                        help='分析标准')
    parser.add_argument('--stability_weight', type=float, default=0.3, help='稳定性权重')
    parser.add_argument('--momentum_weight', type=float, default=0.4, help='动量权重')
    parser.add_argument('--volume_weight', type=float, default=0.2, help='成交量权重')
    parser.add_argument('--frequency_weight', type=float, default=0.1, help='频率权重')
    parser.add_argument('--pe_weight', type=float, default=0.0, help='市盈率权重')
    parser.add_argument('--pb_weight', type=float, default=0.0, help='市净率权重')
    parser.add_argument('--plot', action='store_true', help='是否绘制图表')
    parser.add_argument('--stock_code', type=str, help='要分析的特定股票代码')
    return parser.parse_args()

def main():
    # 解析命令行参数
    args = parse_args()
    
    # 创建分析器实例
    analyzer = StockAnalyzer(args.data_dir)
    
    # 设置权重
    weights = {
        'stability': args.stability_weight,
        'momentum': args.momentum_weight,
        'volume': args.volume_weight,
        'frequency': args.frequency_weight,
        'pe': args.pe_weight,
        'pb': args.pb_weight
    }
    
    # 如果指定了特定股票代码，则分析该股票
    if args.stock_code:
        if args.stock_code in analyzer.stock_info:
            info = analyzer.stock_info[args.stock_code]
            print(f"\n股票信息: {info['名称']}({args.stock_code})")
            print(f"最新价格: {info['数据'][-1]['价格']}元")
            print(f"平均涨跌幅: {info['平均涨跌幅']:.2f}%")
            print(f"60日涨跌幅: {info['60日涨跌幅']:.2f}%")
            print(f"年初至今涨跌幅: {info['年初至今涨跌幅']:.2f}%")
            print(f"市盈率: {info['市盈率']}")
            print(f"市净率: {info['市净率']}")
            print(f"平均换手率: {info['平均换手率']:.2f}%")
            print(f"出现次数: {info['出现次数']}")
            print("\n推荐理由:")
            print(analyzer.get_stock_recommendation_reasons(args.stock_code))
            
            # 如果需要绘图，则绘制该股票的价格走势
            if args.plot:
                analyzer.plot_stock_price_trend(args.stock_code)
        else:
            print(f"找不到股票代码 {args.stock_code} 的数据")
        return
    
    # 根据指定的标准分析股票
    print(f"\n按 {args.criteria} 标准分析股票:")
    analyzer.analyze_by_criteria(args.criteria, args.top_n, args.min_freq, weights if args.criteria == 'comprehensive' else None)
    
    # 生成推荐报告
    print("\n生成推荐股票报告...")
    analyzer.generate_recommendation_report(args.top_n, args.min_freq, weights, args.output)
    
    # 如果需要绘图，则绘制前N只股票的比较图
    if args.plot:
        analyzer.plot_top_stocks_comparison(args.top_n, args.min_freq, weights)

if __name__ == "__main__":
    main()